from bst import BinTreeNode ,BST
class AVLNode(BinTreeNode):
    def __init__(self,data) :
        BinTreeNode.__init__(self,data)
        self.bf=0
class AVLTree(BST):
    def __init__(self,li=None):
        BST.__init__(self,li)
    #左旋
    def rotate_left (self,p,c):# 目前仅适用于插入
        s2=c.lchild
        p.rchild=s2
        if s2:
            s2.parent=p
        c.lchild=p
        p.parent=c

        p.bf=0
        c.bf=0
        return c
    #右旋
    def rorate_right(self,p,c):
        s2=c.rchild
        p.lchild=s2
        if s2 :
            s2.parent=p
        c.rchild=p
        p.parent=c

        p.bf=0
        c.bf=0
        return c
        #右旋——左旋
    def rotate_right_lef(seld,p,c):
        g=c.lchild
        s3=g.rchild
        c.lchid=s3
        if s3:
            s3.parent=c
        g.rchild=c
        c.parent=g

        s2=g.lchild
        p.rchild=s2
        if s2:
            s2.parent=p
        g.lchild=p
        p.parent=g

        #更新bf
        if g.bf>0:
            p.bf=-1
            c.bf=0
        elif g.bf<0:
            p.bf=0
            c.bf=1
        else:#插入的是g
            p.bf=0
            c.bf=0
        g.bf=0
        return g
    def rotate_left_right(self,p,c):
        g=c.rchild

        s2=g.lchild
        c.rchild=s2
        if s2:
            s2.parent=c
        g.lchild=c
        c.parent=g

        s3=g.rchild
        p.lchild=s3
        if s3:
            s3.parent=p
        g.rchild=p
        p.parent=g

        #更新bf
        if g.bf<0:
            c.bf=1
            p.bf=0
        elif g.bf>0:
            p.bf=0
            c.bf=-1
        else:
            p.bf=0
            c.bf=0
        g.bf=0
        return g

    def insert_no_res (self,val):
        #1.插入
        p=self.root
        if not p:#空树
            self.root=AVLNode(val)
            return
        while True:
            if val<p.data:
                if p.lchild:
                    p=p.lchild
                else:
                    p.lchild=AVLNode(val)
                    p.lchild.parent=p
                    node=p.lchild#node 存储的就是插入的节点
                    break
            elif val>p.data:
                if p.rchild:
                    p=p.rchild
                else:
                    p.rchild=AVLNode(val)
                    p.rchild.parent=p
                    node=p.rchild
                    break
            else:#val==p.data
                return
        #2.更新balance factor
        while node.parent:#node.parent不空
            if node.parent.lchild==node:
                if node.parent.bf<0:# 原来node.parent.bf==-1,左子树更沉了
                    #看node哪边沉
                    g=node.parent.parent#为了连接旋转之后的子树
                    x=node.parent#旋转前的子树的根
                    if node.bf>0:
                        self.rotate_left_right(node.parent,node)
                    else:
                        n=self.rorate_right(node.parent,node)
                            #记得把n连接起来
                        
                elif node.parent.bf>0:# 原来node.parent.bf=-1，更新后变为0
                    node.parent.bf=0
                    break
                else:#原来的node.parent.bf=0更新后变为-1
                    node.parent=-1
                    node=node.parent
                    continue

            else:#传递是从右子树来的，右子树更沉了
                #更新后node.parent+=1
                if node.parent.bf>0:#说明原来ndoe.parnet.bf==1,更新后变为2
                    #做旋转
                    #看node.parent==2
                    g=node.parent.parent
                    x=node.parent#旋转前子树
                    if node.bf<0:
                        n=self.rotate_right_lef(node.parent,node)
                    else:
                        n=self.rotate_left(node.parent,node)
                    #记得连起来
                elif node.parent.bf<0:#说明原来ndoe.parnet.bf==-1,更新后变为0
                    node.parent.bf=0
                    break
                else:
                    node.parent.bf=1
                    node=node.parent
                    continue
            #连接旋转后的子树
            n.parent=g
            if g:#如果g不是空
                if x ==g.lchild:
                    g.lchild = n
                else:
                    g.rchild=n
                break
            else:
                self.root=n
                break
if __name__=="__main__":
    #tree=AVLTree([9,8,7,6,5,4,3,2,1])
    print("xx")
    tree=AVLTree([9,8,7,6,5,4,3,2,1])
    print("hh")
    tree.pre_order(tree.root)
    print("")
    tree.in_order(tree.root)

    
